import os
import json
from tools import File
from tools.chat import Survey

init_data = {
    "2006": [],
    "2007": [],
    "2008": [],
    "2009": [],
    "2010": [],
    "2011": [],
    "2012": [],
    "2013": [],
    "2014": [],
    "2015": [],
    "2016": [],
    "2017": [],
    "2018": [],
    "2019": [],
    "2020": [],
    "2021": [],
    "2022": [],
    "2023": [],
}

# Create output directory if it doesn't exist
os.makedirs("output", exist_ok=True)

folder_path = "./test"

file_handler = File()
file_paths = file_handler.read_folder(folder_path)
survey = Survey()


def survey_reports(llms):
    for llm in llms:
        survey_report(llm)


def survey_report(llm):
    for file_name, year, file_path in file_paths:
        output_file_path = f"./output/{llm}_output.json"
        os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
        if os.path.exists(output_file_path):
            with open(output_file_path, "r", encoding="utf-8") as file:
                try:
                    data = json.load(file)
                except json.JSONDecodeError:
                    print("文件内容损坏!")
        else:
            with open(output_file_path, "w", encoding="utf-8") as file:
                json.dump(init_data, file, ensure_ascii=False, indent=4)
            data = init_data

        try:
            context = file_handler.search_text(file_path)
            content_rate = survey.rate(llm, {"text": context})
            if content_rate is None:
                score = {
                    "artificial_intelligence_application": False,
                    "natural_language_processing": False,
                    "machine_learning": False,
                    "neural_network": False,
                    "automatic_driving": False,
                    "generative_artificial_intelligence": False,
                    "visual_image_recognition": False,
                    "robot_technology": False,
                }
            else:
                score = dict(content_rate)
            print("*" * 50)
            print(dict(score))
            info = {file_name: score}
            data[year].append(info)
            file_handler.save_json(f"./output/{llm}_output.json", data)
        except Exception as e:
            print(e)


if __name__ == "__main__":
    llms = ["glm4", "qwen", "llama", "gemma"]
    survey_reports(llms)
